#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy, sys
import moveit_commander

import tf.transformations as tfs
import numpy as np
import math
import kinematics as ks

from arm_manager import ArmManager
from gripper_manager import GripperManager
from chassis import ChassisManager
from sensors_init import SensorsInit

from moveit_msgs.msg import RobotTrajectory
from trajectory_msgs.msg import JointTrajectoryPoint
from geometry_msgs.msg import PoseStamped, Pose, WrenchStamped
from std_msgs.msg import UInt8
from std_msgs.msg import Float64MultiArray
from webots_ros.msg import RecognitionObject
from webots_ros.msg import Float64Stamped


class TeachPendant:
    def __init__(self):
        # 初始化ROS节点
        rospy.init_node('teach_pendant_demo')
        
        #arm
        self.cmarm = ArmManager()
        #gripper
        self.cmgripper = GripperManager()
        #chassis
        # self.cm = ChassisManager()
        #camera
        self.si = SensorsInit()
        self.si.enable_sensors()
    
    def threshold(self,fvalue):
        thd = 3   #  N
        jvalue = fvalue* np.sign(fvalue)- thd
        if(jvalue>0):
            return np.sign(fvalue)*jvalue
        else:
            return 0 

    def run(self):
        Dd = 1
        Ad = 1 
        Ts = 0.03
        xd_k_1_dot = np.zeros((3, 1))
        while True:
        # for i in range(100):
            try:
                joint_0_c = rospy.wait_for_message("/robot/joint1_sensor/value", Float64Stamped, timeout=10)
                joint_1_c = rospy.wait_for_message("/robot/joint2_sensor/value", Float64Stamped, timeout=10)
                joint_2_c = rospy.wait_for_message("/robot/joint4_sensor/value", Float64Stamped, timeout=10)
                Fext_msg  = rospy.wait_for_message("/robot/touch_sensor/values", WrenchStamped, timeout=10)
            except:
                pass

            pose_c = ks.FK(joint_0_c.data,joint_1_c.data,joint_2_c.data)
            
            force_x = Fext_msg.wrench.force.x
            force_z = Fext_msg.wrench.force.z
            
            force_x_cartesian = -(math.cos(pose_c[2])*force_z + math.sin(pose_c[2])*force_x)
            force_y_cartesian = math.cos(pose_c[2])*force_x - math.sin(pose_c[2])*force_z

            force_x_cartesian = self.threshold(force_x_cartesian)
            force_y_cartesian = self.threshold(force_y_cartesian)

            Fext = [[force_x_cartesian],[force_y_cartesian],[0]]
            
            Jacb_inv = ks.Jacobian(joint_1_c.data+ joint_2_c.data- math.pi/2, joint_1_c.data, inverse=True)

            # xd_k_1 = ks.FK(joint_0_c.data,joint_1_c.data,joint_2_c.data)
            # xd_k_1 = xd_k_1.reshape((3,1))

            xd_k_dot = ((Fext- Dd* xd_k_1_dot)/ Ad*  Ts+ xd_k_1_dot)/2.0

            qd_dot = np.dot(Jacb_inv, xd_k_dot)
            # 指令下发
            self.cmarm.set_velocity_for_each_joint(qd_dot.reshape((3)))

            # 更新
            xd_k_1_dot = xd_k_dot 

            print(qd_dot)
            rospy.sleep(Ts)

if __name__ == "__main__":
    # try:
    #     TeachPendant()
    # except rospy.ROSInterruptException:
    #     pass
    TP = TeachPendant()
    TP.run()


# y = 0.6

# theta0 = 0.7
# l0 = 0.112
# l1 = 0.6035+0.0415
# l2 = 0.1175
# l3 = 0.14
# f1 = 0.9323
# PI = 3.141592
# theta1 = math.asin((y+theta0-f1)/l1) + PI/2
# theta3 = PI/2 - theta1 

# x = math.cos(theta1-PI/2)*l1 + l2 + l0

# print(theta1)
# print(theta3)
# print(x)